import numpy as np
import libpysal as lps
from mgwr.gwr import GWR, MGWR
from mgwr.sel_bw import Sel_BW
from mgwr.utils import shift_colormap, truncate_colormap
import geopandas as gpd
import matplotlib.pyplot as plt
import matplotlib as mpl
import pandas as pd

plt.rcParams['font.sans-serif'] = ['KaiTi']
plt.rcParams['axes.unicode_minus'] = False


def gwr_mgwr(geo_data: object, y: object, x_list: object) -> object:
    g_y = geo_data[y].values.reshape((-1, 1))
    g_X = geo_data[x_list].values
    u = geo_data['lon']
    v = geo_data['lat']
    g_coords = list(zip(u, v))
    g_X = (g_X - g_X.mean(axis=0)) / g_X.std(axis=0)
    g_y = g_y.reshape((-1, 1))
    g_y = (g_y - g_y.mean(axis=0)) / g_y.std(axis=0)
    # Calibrate GWR model

    gwr_selector = Sel_BW(g_coords, g_y, g_X)
    gwr_bw = gwr_selector.search(bw_min=2)
    print(gwr_bw)
    gwr_results = GWR(g_coords, g_y, g_X, gwr_bw).fit()

    # Calibrate MGWR model

    mgwr_selector = Sel_BW(g_coords, g_y, g_X, multi=True)
    mgwr_bw = mgwr_selector.search(multi_bw_min=[2])
    print(mgwr_bw)
    mgwr_results = MGWR(g_coords, g_y, g_X, mgwr_selector).fit()

    print(gwr_results.summary())

    print(mgwr_results.summary())

    # Prepare GWR results for mapping
    # Add GWR parameters to GeoDataframe
    geo_data['gwr_intercept'] = gwr_results.params[:, 0]
    count = 1
    for i in x_list:
        geo_data['gwr_' + i] = gwr_results.params[:, count]
        count = count + 1

    # Obtain t-vals filtered based on multiple testing correction
    gwr_filtered_t = gwr_results.filter_tvals()

    # Prepare MGWR results for mapping

    # Add MGWR parameters to GeoDataframe
    geo_data['mgwr_intercept'] = mgwr_results.params[:, 0]
    count = 1
    for i in x_list:
        geo_data['mgwr_' + i] = mgwr_results.params[:, count]
        count = count + 1

    # Obtain t-vals filtered based on multiple testing correction
    mgwr_filtered_t = mgwr_results.filter_tvals()

    # Comparison maps of GWR vs. MGWR parameter surfaces where the grey units pertain to statistically insignificant parameters
    # Comparison maps of GWR vs. MGWR parameter surfaces where the grey units pertain to statistically insignificant parameters

    # Prep plot and add axes
    fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(45, 20))
    ax0 = axes[0]
    ax0.set_title('Intercept GWR (BW: ' + str(gwr_bw) + ')', fontsize=40)
    ax1 = axes[1]
    ax1.set_title('Intercept MGWR (BW: ' + str(mgwr_bw[0]) + ')', fontsize=40)

    # Set color map
    cmap = plt.cm.seismic

    # Find min and max values of the two combined datasets
    gwr_min = geo_data['gwr_intercept'].min()
    gwr_max = geo_data['gwr_intercept'].max()
    mgwr_min = geo_data['mgwr_intercept'].min()
    mgwr_max = geo_data['mgwr_intercept'].max()
    vmin = np.min([gwr_min, mgwr_min])
    vmax = np.max([gwr_max, mgwr_max])
    # If all values are negative use the negative half of the colormap
    if (vmin < 0) & (vmax < 0):
        cmap = truncate_colormap(cmap, 0.0, 0.5)
    # If all values are positive use the positive half of the colormap
    elif (vmin > 0) & (vmax > 0):
        cmap = truncate_colormap(cmap, 0.5, 1.0)
    # Otherwise, there are positive and negative values so the colormap so zero is the midpoint
    else:
        cmap = shift_colormap(cmap, start=0.0, midpoint=1 - vmax / (vmax + abs(vmin)), stop=1.)

    # Create scalar mappable for colorbar and stretch colormap across range of data values
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=vmin, vmax=vmax))

    # Plot GWR parameters
    geo_data.plot('gwr_intercept', cmap=sm.cmap, ax=ax0, vmin=vmin, vmax=vmax,
                  **{'edgecolor': 'black', 'alpha': .65})
    # If there are insignificnt parameters plot gray polygons over them
    if (gwr_filtered_t[:, 0] == 0).any():
        geo_data[gwr_filtered_t[:, 0] == 0].plot(color='lightgrey', ax=ax0, **{'edgecolor': 'black'})

    # Plot MGWR parameters
    geo_data.plot('mgwr_intercept', cmap=sm.cmap, ax=ax1, vmin=vmin, vmax=vmax,
                  **{'edgecolor': 'black', 'alpha': .65})
    # If there are insignificnt parameters plot gray polygons over them
    if (mgwr_filtered_t[:, 0] == 0).any():
        geo_data[mgwr_filtered_t[:, 0] == 0].plot(color='lightgrey', ax=ax1, **{'edgecolor': 'black'})
    # Set figure options and plot
    fig.tight_layout()
    fig.subplots_adjust(right=0.9)
    cax = fig.add_axes([0.92, 0.14, 0.03, 0.75])
    sm._A = []
    cbar = fig.colorbar(sm, cax=cax)
    cbar.ax.tick_params(labelsize=50)
    ax0.get_xaxis().set_visible(False)
    ax0.get_yaxis().set_visible(False)
    ax1.get_xaxis().set_visible(False)
    ax1.get_yaxis().set_visible(False)
    plt.show()
    del fig

    # Prep plot and add axes
    count = 1
    for i in x_list:
        fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(45, 20))
        ax0 = axes[0]
        ax0.set_title(i + 'GWR (BW: ' + str(gwr_bw) + ')', fontsize=40)
        ax1 = axes[1]
        ax1.set_title(i + 'MGWR (BW: ' + str(mgwr_bw[count]) + ')', fontsize=40)

        # Set color map
        cmap = plt.cm.seismic

        # Find min and max values of the two combined datasets
        gwr_min = geo_data['gwr_' + i].min()
        gwr_max = geo_data['gwr_' + i].max()
        mgwr_min = geo_data['mgwr_' + i].min()
        mgwr_max = geo_data['mgwr_' + i].max()
        vmin = np.min([gwr_min, mgwr_min])
        vmax = np.max([gwr_max, mgwr_max])
        # If all values are negative use the negative half of the colormap
        if (vmin < 0) & (vmax < 0):
            cmap = truncate_colormap(cmap, 0.0, 0.5)
        # If all values are positive use the positive half of the colormap
        elif (vmin > 0) & (vmax > 0):
            cmap = truncate_colormap(cmap, 0.5, 1.0)
        # Otherwise, there are positive and negative values so the colormap so zero is the midpoint
        else:
            cmap = shift_colormap(cmap, start=0.0, midpoint=1 - vmax / (vmax + abs(vmin)), stop=1.)

        # Create scalar mappable for colorbar and stretch colormap across range of data values
        sm = plt.cm.ScalarMappable(cmap=cmap, norm=plt.Normalize(vmin=vmin, vmax=vmax))

        # Plot GWR parameters
        geo_data.plot('gwr_intercept', cmap=sm.cmap, ax=ax0, vmin=vmin, vmax=vmax,
                      **{'edgecolor': 'black', 'alpha': .65})
        # If there are insignificnt parameters plot gray polygons over them
        if (gwr_filtered_t[:, count] == 0).any():
            geo_data[gwr_filtered_t[:, count] == 0].plot(color='lightgrey', ax=ax0, **{'edgecolor': 'black'})

        # Plot MGWR parameters
        geo_data.plot('mgwr_intercept', cmap=sm.cmap, ax=ax1, vmin=vmin, vmax=vmax,
                      **{'edgecolor': 'black', 'alpha': .65})
        # If there are insignificnt parameters plot gray polygons over them
        if (mgwr_filtered_t[:, count] == 0).any():
            geo_data[mgwr_filtered_t[:, count] == 0].plot(color='lightgrey', ax=ax1, **{'edgecolor': 'black'})
        # Set figure options and plot
        fig.tight_layout()
        fig.subplots_adjust(right=0.9)
        cax = fig.add_axes([0.92, 0.14, 0.03, 0.75])
        sm._A = []
        cbar = fig.colorbar(sm, cax=cax)
        cbar.ax.tick_params(labelsize=50)
        ax0.get_xaxis().set_visible(False)
        ax0.get_yaxis().set_visible(False)
        ax1.get_xaxis().set_visible(False)
        ax1.get_yaxis().set_visible(False)
        count = count + 1
        plt.show()
        del fig
        return gwr_results.summary(), mgwr_results.summary()


# geo_data = gpd.read_file('C:/Users/18295335197/Desktop/test/街区GWR.geojson')
# y = 'infected_count'
# x_list = ['hospital_count', 'shop_count', 'catering_count']
# gwr_mgwr(geo_data, y, x_list)
